
"""
任务改写
"""
from core.data_analyze.state import AnalyzeState
from langgraph.types import Command
from core.data_analyze.prompts import data_analyze_prompts
from langchain_core.messages import AIMessage
from langgraph.config import get_stream_writer
from core.llm import llm
from loguru import logger
import datetime


def task_rewrite_node(state: AnalyzeState):

    writer = get_stream_writer()
    writer("## **改写任务描述**")
    logger.info("## **改写任务描述**")
    template_prompt = data_analyze_prompts.get('task_rewrite')
    prompt = template_prompt.format(
        inputs={
            "task_description": state.get('task_description'),
            "current_date": str(datetime.date.today())
        },
        remove_template_variables=True
    )
    revised_task = llm.invoke(prompt).content
    writer(f"改写后的任务描述：{revised_task}")
    logger.info(f"改写后的任务描述：{revised_task}")
    return Command(
        update={
            "messages": [AIMessage(content=revised_task, name="task_rewrite")],
            "revised_task": revised_task
        },
        goto="fetch_data",
    )
